#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray

class VelocityController(Node):
    def __init__(self):
        super().__init__('stewart_velocity_controller')
        print("VelocityController initialized")
        # 控制参数
        self.startup_speed = 0.4
        self.normal_speed_s1_s4_max = 0.15  # s1-s4正常运动最大速度
        self.normal_speed_s5_s6 = 0.1       # s5-s6正常运动速度
        self.Kp = 1.0
        self.Ki = 0.05
        self.Kd = 0.1
        self.max_velocity = 0.4
        self.min_velocity = 0.3
        self.deadzone_threshold = 0.005  # 死区阈值 (5mm)
        self.integral_reset_threshold = 0.01

        self.integral_errors = [0.0] * 6
        self.integral_limit = 0.5

        self.joint_names = ['s1', 's2', 's3', 's4', 's5', 's6']
        self.num_joints = len(self.joint_names)

        self.target_lengths = [0.0] * self.num_joints
        self.current_lengths = [0.0] * self.num_joints
        self.current_velocities = [0.0] * self.num_joints

        # 启动前的长度记录
        self.startup_lengths = [0.0] * self.num_joints
        self.in_startup = [True] * self.num_joints

        self.joint_state_sub = self.create_subscription(
            JointState, '/joint_states', self.joint_state_callback, 10)

        self.target_sub = self.create_subscription(
            Float64MultiArray, '/target_lengths', 
            self.target_lengths_callback, 10)

        self.velocity_pub = self.create_publisher(
            Float64MultiArray, '/rod_velocity_controller/commands', 10)

        self.timer = self.create_timer(0.05, self.control_loop)
    
    def target_lengths_callback(self, msg):
        if len(msg.data) == 6:
            self.target_lengths = msg.data
        else:
            self.get_logger().warn('Received invalid target lengths. Expected 6 elements')

    def joint_state_callback(self, msg: JointState):
        if all(pos == 0 for pos in msg.position) and all(vel == 0 for vel in msg.velocity):
            return
        for i, name in enumerate(self.joint_names):
            if name in msg.name:
                idx = msg.name.index(name)
                self.current_lengths[i] = msg.position[idx]
                if len(msg.velocity) > idx:
                    self.current_velocities[i] = msg.velocity[idx]

    def control_loop(self):
        velocity_commands = []
        current_positions = self.current_lengths

        # 启动阶段判断与切换
        for i in range(self.num_joints):
            if self.in_startup[i]:
                # 记录启动前长度
                self.startup_lengths[i] = self.current_lengths[i]
                self.in_startup[i] = False

        for i in range(self.num_joints):
            error = self.target_lengths[i] - self.current_lengths[i]

            # 判断是否切换为正常运动速度
            if abs(self.current_lengths[i] - self.startup_lengths[i]) > 0.05:
                self.in_startup[i] = False

            # PID控制
            if abs(error) < self.deadzone_threshold:
                velocity = 0.0
                self.integral_errors[i] *= 0.5
            else:
                self.integral_errors[i] += error * 0.05
                self.integral_errors[i] = max(min(self.integral_errors[i], self.integral_limit), -self.integral_limit)
                d_error = -self.current_velocities[i]
                velocity = (self.Kp * error + self.Ki * self.integral_errors[i] + self.Kd * d_error)

            # 死区处理：允许速度在[-max,-0.05]和[0.05,max]之间
            if 0 < abs(velocity) < self.min_velocity:
                if velocity > 0:
                    velocity = max(velocity, 0.05)
                else:
                    velocity = min(velocity, -0.05)

            # 限制最大速度
            if i < 4:
                max_v = self.normal_speed_s1_s4_max
            else:
                max_v = self.normal_speed_s5_s6
            velocity = max(min(velocity, max_v), -max_v)

            # 启动阶段速度
            if not self.in_startup[i]:
                # 正常运动速度
                pass  # 已经限制
            else:
                # 启动速度
                velocity = self.startup_speed * (1 if error > 0 else -1)

            velocity_commands.append(velocity)

        msg = Float64MultiArray()
        msg.data = velocity_commands
        self.velocity_pub.publish(msg)

def main(args=None):
    rclpy.init(args=args)
    node = VelocityController()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()